#!/usr/bin/env python3

import numpy as np
from PIL import Image
import os
import cv2

# ==============================================================================
# HELPER AND UTILITY FUNCTIONS
# ==============================================================================

def imread_cv2(path, options=cv2.IMREAD_COLOR):
    """Open an image or a depthmap with opencv-python."""
    if path.endswith((".exr", "EXR")):
        options = cv2.IMREAD_ANYDEPTH
    img = cv2.imread(path, options)
    if img is None:
        raise IOError(f"Could not load image={path} with {options=}")
    if img.ndim == 3:
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    return img

# ==============================================================================
# UNPROJECTION LOGIC
# ==============================================================================

def depthmap_to_camera_coordinates(depthmap, camera_intrinsics):
    """
    Unprojects a depth map into a 3D point cloud in the camera's coordinate system.
    """
    camera_intrinsics = np.float32(camera_intrinsics)
    H, W = depthmap.shape

    assert camera_intrinsics[0, 1] == 0.0 and camera_intrinsics[1, 0] == 0.0
    fu, fv = camera_intrinsics[0, 0], camera_intrinsics[1, 1]
    cu, cv = camera_intrinsics[0, 2], camera_intrinsics[1, 2]

    u, v = np.meshgrid(np.arange(W), np.arange(H))
    z_cam = depthmap
    x_cam = (u - cu) * z_cam / fu
    y_cam = (v - cv) * z_cam / fv
    X_cam = np.stack((x_cam, y_cam, z_cam), axis=-1).astype(np.float32)

    valid_mask = depthmap > 0.0
    return X_cam, valid_mask


def depthmap_to_absolute_camera_coordinates(depthmap, camera_intrinsics, camera_pose):
    """
    Unprojects a depth map to a 3D point cloud in absolute world coordinates.
    """
    X_cam, valid_mask = depthmap_to_camera_coordinates(depthmap, camera_intrinsics)

    R_cam2world = camera_pose[:3, :3]
    t_cam2world = camera_pose[:3, 3]

    # Transform points from camera coordinates to world coordinates
    X_world = np.einsum("ik,hwk->hwi", R_cam2world, X_cam) + t_cam2world[None, None, :]
    
    return X_world, valid_mask

# ==============================================================================
# ROTATION LOGIC (Both Versions)
# ==============================================================================

# --- VGGT VERSION ---
def rotate_90_degrees(image, depth_map, extri_opencv, intri_opencv, clockwise=True):
    image_height, image_width = image.shape[:2]
    # Image/Depth rotation
    if clockwise:
        k = -1
    else:
        k = 1
    rotated_image = np.rot90(image, k)
    rotated_depth_map = np.rot90(depth_map, k) if depth_map is not None else None
    
    # Intrinsic adjustment
    fx, fy, cx, cy = intri_opencv[0,0], intri_opencv[1,1], intri_opencv[0,2], intri_opencv[1,2]
    new_intri_opencv = np.eye(3, dtype=intri_opencv.dtype)
    new_intri_opencv[0,0], new_intri_opencv[1,1] = fy, fx
    if clockwise:
        new_intri_opencv[0,2], new_intri_opencv[1,2] = image_height - cy - 1, cx
    else:
        new_intri_opencv[0,2], new_intri_opencv[1,2] = cy, image_width - cx - 1
        
    # Extrinsic adjustment
    R, t = extri_opencv[:, :3], extri_opencv[:, 3]
    if not clockwise:
        R_rot = np.array([[0, -1, 0], [1, 0, 0], [0, 0, 1]])
    else:
        R_rot = np.array([[0, 1, 0], [-1, 0, 0], [0, 0, 1]])
    new_R = R @ R_rot
    new_t = t @ R_rot
    new_extri_opencv = np.hstack((new_R, new_t.reshape(-1, 1)))
    return rotated_image, rotated_depth_map, new_extri_opencv, new_intri_opencv

# --- ANOTHER VERSION ---
def rotate(image, depth, intrinsics, pose, k=1):
    H, W, _ = image.shape
    image_rot, depth_rot = np.rot90(image, k), np.rot90(depth, k)
    
    if k % 4 == 1: # CCW
        
        R_transform = np.array([[0,-1,0,0],[1,0,0,0],[0,0,1,0],[0,0,0,1]], dtype=np.float32)
    elif k % 4 == -1 or k % 4 == 3: # CW
        R_transform = np.array([[0,1,0,0],[-1,0,0,0],[0,0,1,0],[0,0,0,1]], dtype=np.float32)
    else:
        raise NotImplementedError
    # new_pose = R_transform @ pose
    new_pose = pose @ R_transform
    
    fx, fy, cx, cy = intrinsics[0,0], intrinsics[1,1], intrinsics[0,2], intrinsics[1,2]
    fx_new, fy_new = fy, fx
    if k % 4 == 1:
        cx_new, cy_new = cy, W - cx - 1
    else:
        cx_new, cy_new = H - cy - 1, cx
    intrinsics_new = np.array([[fx_new,0,cx_new],[0,fy_new,cy_new],[0,0,1]], dtype=np.float32)
    return image_rot, depth_rot, intrinsics_new, new_pose

# ==============================================================================
# MAIN VERIFICATION SCRIPT
# ==============================================================================

def get_colored_point_cloud(image, depth, intrinsics, pose):
    """Helper to unproject data and return a sorted, colored point cloud."""
    # Unproject to get 3D points in world coordinates
    points_world, valid_mask = depthmap_to_absolute_camera_coordinates(depth, intrinsics, pose)
    
    # Get valid points and their corresponding colors
    valid_points = points_world[valid_mask]
    valid_colors = image[valid_mask]
    
    # Sort the points lexicographically (by z, then y, then x) for consistent comparison
    # This ensures that even if pixel order changes, the set of points can be compared.
    indices = np.lexsort((valid_points[:, 0], valid_points[:, 1], valid_points[:, 2]))
    
    return valid_points[indices], valid_colors[indices]

def main():
    """Main function to load data and run verifications."""
    base_path = '/home/liucong/data/3d/hypersim_processed/train/ai_019_007/cam_00'
    rgb_path = os.path.join(base_path, '000008_rgb.png')
    depth_path = os.path.join(base_path, '000008_depth.npy')
    cam_path = os.path.join(base_path, '000008_cam.npz')

    # Load all data
    original_image = imread_cv2(rgb_path)
    original_depth = np.load(depth_path).astype(np.float32)
    original_depth[~np.isfinite(original_depth)] = 0
    cam_file = np.load(cam_path)
    original_intri = cam_file["intrinsics"].astype(np.float32)
    original_pose = cam_file["pose"].astype(np.float32)

    print("--- Verification of 3D Point Cloud Invariance After Rotation ---")

    # 1. Get the original colored point cloud
    print("\n1. Generating original point cloud...")
    points_orig, colors_orig = get_colored_point_cloud(
        original_image, original_depth, original_intri, original_pose
    )
    print(f"   ...found {len(points_orig)} valid points.")

    # 2. Test the VGGT version
    print("\n2. Testing VGGT version (CCW rotation)...")
    vggt_extri_3x4 = original_pose[:3, :] # VGGT version needs a 3x4 matrix
    (img_vggt, depth_vggt, extri_vggt, intri_vggt) = rotate_90_degrees(
        original_image, original_depth, vggt_extri_3x4, original_intri, clockwise=False
    )
    pose_vggt = np.vstack([extri_vggt, [0, 0, 0, 1]]) # Rebuild 4x4 pose for unprojection
    points_vggt, colors_vggt = get_colored_point_cloud(
        img_vggt, depth_vggt, intri_vggt, pose_vggt
    )
    print(f"   ...found {len(points_vggt)} valid points after rotation.")
    print(f"   ✅ 3D points are identical: {np.allclose(points_orig, points_vggt)}")
    print(f"   ✅ Colors are identical: {np.array_equal(colors_orig, colors_vggt)}")

    # 3. Test the 'another' version
    print("\n3. Testing 'another' version (CCW rotation)...")
    (img_other, depth_other, intri_other, pose_other) = rotate(
        original_image, original_depth, original_intri, original_pose, k=-1
    )
    points_other, colors_other = get_colored_point_cloud(
        img_other, depth_other, intri_other, pose_other
    )
    print(f"   ...found {len(points_other)} valid points after rotation.")
    print(f"   ✅ 3D points are identical: {np.allclose(points_orig, points_other)}")
    print(f"   ✅ Colors are identical: {np.array_equal(colors_orig, colors_other)}")
    print("-" * 40)


if __name__ == '__main__':
    try:
        main()
    except FileNotFoundError:
        print("\n❌ ERROR: Data files not found.")
        print("Please ensure your data files are in the same directory as the script.")
    except KeyError as e:
        print(f"\n❌ ERROR: A key was not found in the camera parameters file: {e}")
        print("Please ensure '000008_cam.npz' contains 'intrinsics' and 'pose'.")